/**
 * 
 */
package edu.umd.clip.lm.storage;

import java.text.DecimalFormat;
import java.text.NumberFormat;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.*;
import java.util.concurrent.locks.*;
import java.util.*;

import edu.umd.clip.lm.model.OnDiskCompactProbTree;
import edu.umd.clip.lm.model.ProbTree;
import edu.umd.clip.lm.util.*;
import edu.umd.clip.lm.storage.AbstractProbTreeStorage.Key;
import edu.umd.clip.lm.storage.remote.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class ClientProbTreeStorage extends AbstractProbTreeStorage implements ClientListener {
	private ClientStorageConnection connection;
	private ConcurrentHashMap<Integer, Message> sentMessages;
	private static final int NUM_SEGMENTS = 16;
	private static AtomicLong receivedHits = new AtomicLong();
	private static AtomicLong requestedHits = new AtomicLong();
	private static AtomicLong lostRequests = new AtomicLong();
	private static AtomicLong timeWaited = new AtomicLong();
	private static AtomicLong asyncRequests = new AtomicLong();
	private Segment segments[];
	
	private static class Segment {
		//private LRU<Key, ProbTree> receivedProbTrees;
		//private HashMap<Key,Key> requestedKeys;
		private Lock lock;
		private HashMap<Key,List<AsyncConsumer<Key,OnDiskCompactProbTree>>> consumers;

		private Segment() {
			//receivedProbTrees = new LRU<Key,ProbTree>(MAX_PROB_TREES_WAITING / NUM_SEGMENTS);
			//requestedKeys = new HashMap<Key,Key>(64);
			
			lock = new ReentrantLock();
			//receivedSomething = lock.newCondition(); 
			consumers = new HashMap<Key,List<AsyncConsumer<Key,OnDiskCompactProbTree>>>(); 
		}
	}
	
	public ClientProbTreeStorage(ClientStorageConnection connection) {
		this.connection = connection;
		segments = new Segment[NUM_SEGMENTS];
		for(int i=0; i<NUM_SEGMENTS; ++i) {
			segments[i] = new Segment();
		}
		sentMessages = new ConcurrentHashMap<Integer,Message>(64);
		connection.addListener(this);
	}

	private Segment getSegment(Key key) {
		int h = key.hashCode();
        h ^= (h >>> 21) ^ (h >>> 13);
        h ^= (h >>> 8) ^ (h >>> 5);

        h &= NUM_SEGMENTS - 1;
        return segments[h];
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.storage.AbstractProbTreeStorage#getProbTree(edu.umd.clip.lm.storage.AbstractProbTreeStorage.Key)
	 */
	@Override
	public OnDiskCompactProbTree getProbTree(Key key) {
		long time = 0;
		final AtomicBoolean isDone = new AtomicBoolean(false);
		final AtomicReference<OnDiskCompactProbTree> probTree = new AtomicReference<OnDiskCompactProbTree>(null);
		
		AsyncConsumer<Key,OnDiskCompactProbTree> consumer = new AsyncConsumer<Key,OnDiskCompactProbTree>() {
			@Override
			public synchronized void receivedData(Key key, OnDiskCompactProbTree value) {
				probTree.set(value); 
				isDone.set(true);
				notify();
			}
			
		};
		request(key, consumer);
		
		synchronized(consumer) {
			while(!isDone.get()) {
				long start = System.nanoTime();
				try {
					consumer.wait();
				} catch (InterruptedException e) {
					e.printStackTrace();
					return null;
				}
				long end = System.nanoTime();
				time += end - start;
			}
		}

		if (time != 0) {
			timeWaited.addAndGet(time);
		}

		return probTree.get();
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.storage.AbstractProbTreeStorage#getStats()
	 */
	@Override
	public String getStats() {
		NumberFormat nf = new DecimalFormat("#.##");

		StringBuilder sb = new StringBuilder();
		long numRequests = connection.getMessagesSent(); 
		sb.append(String.format("ClientProbTree: %d requests, %sms response delay, %d syncs, %sms avg",
				numRequests,
				nf.format(connection.getResponseProcessingTime() / numRequests * 1e-6),
				RequestBundle.getSyncCount(),
				nf.format(RequestBundle.getSyncTime() / RequestBundle.getSyncCount() * 1e-6)
				));
		return sb.toString();
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.storage.AbstractProbTreeStorage#requestProbTree(edu.umd.clip.lm.storage.AbstractProbTreeStorage.Key)
	 */
	@Override
	protected void requestProbTree(Key key) {
		throw new UnsupportedOperationException();
	}

	private Message makeRequest(Segment seg, Key key) {
		Message message = new SingleTreeRequestMessage(key.modelid, key.clusterid, key.factorBits);
		sentMessages.put(message.getId(), message);
		return message;
		//seg.requestedKeys.put(key,key);
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.storage.AbstractProbTreeStorage#setProbTree(edu.umd.clip.lm.storage.AbstractProbTreeStorage.Key, edu.umd.clip.lm.model.ProbTree)
	 */
	@Override
	public void setProbTree(Key key, OnDiskCompactProbTree probTree) {
		throw new UnsupportedOperationException();
	}

	public void receivedReply(Message message) {
		
		SingleTreeRequestMessage originalMessage = (SingleTreeRequestMessage) sentMessages.remove(message.getReplyTo());
		if (originalMessage == null) {
			System.err.printf("received response to unknown message %d\n", message.getReplyTo());
		} else {
			OnDiskCompactProbTree probTree = null;
			if (message instanceof SingleTreeResponseMessage) {
				SingleTreeResponseMessage msg = (SingleTreeResponseMessage) message;
				//seg.receivedProbTrees.put(key, msg.getProbTree());
				probTree = msg.getProbTree();
			} else if (message instanceof FailureResponseMessage) {
				FailureResponseMessage msg = (FailureResponseMessage) message;
				if (msg.getCode() == FailureResponseMessage.NOT_FOUND) {
					//seg.receivedProbTrees.put(key, null);
				} else {
					System.err.printf("unknown error message received\n");
				}
			}

			Key key = new Key(originalMessage.getModelId(), originalMessage.getClusterid(), originalMessage.getFactorBits());
			Segment seg = getSegment(key);
			seg.lock.lock();
			List<AsyncConsumer<Key,OnDiskCompactProbTree>> interestedConsumers = seg.consumers.remove(key);
			seg.lock.unlock();
			
			if (interestedConsumers != null) {
				for(AsyncConsumer<Key,OnDiskCompactProbTree> consumer : interestedConsumers) {
					consumer.receivedData(key, probTree);
				}
			}
		}
	}
	
	public void reconnected() {
		for(Segment seg : segments) {
			seg.lock.lock();
		}
		sentMessages.clear();
		for(Segment seg : segments) {
			if (!seg.consumers.isEmpty()) {
				// resend all outstanding messages
				Key keys[] = seg.consumers.keySet().toArray(new Key[0]);
				
				for(Key key : keys) {
					Message message = makeRequest(seg, key);
					connection.sendMessage(message);
				}
			}
		}
		for(Segment seg : segments) {
			seg.lock.unlock();
		}
	}
	
	public static long getReceivedHits() {
		return receivedHits.longValue();
	}
	
	public static long getRequestedHits() {
		return requestedHits.longValue();
	}
	
	public static long getLostRequests() {
		return lostRequests.longValue();
	}
	
	public static long getTimeWaited() {
		return timeWaited.longValue();
	}
	
	public static long getAsyncRequests() {
		return asyncRequests.longValue();
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.storage.AsyncProducer#request(java.lang.Object, edu.umd.clip.lm.storage.AsyncConsumer)
	 */
	@Override
	public void request(Key key, AsyncConsumer<Key, OnDiskCompactProbTree> consumer) {
		asyncRequests.incrementAndGet();
		Segment seg = getSegment(key);
		seg.lock.lock();
		
		List<AsyncConsumer<Key,OnDiskCompactProbTree>> interestedConsumers = seg.consumers.get(key);
		Message message = null;
		if (interestedConsumers == null) {
			interestedConsumers = new ArrayList<AsyncConsumer<Key,OnDiskCompactProbTree>>(3);
			seg.consumers.put(key, interestedConsumers);
			message = makeRequest(seg, key);
		}
		interestedConsumers.add(consumer);			
		seg.lock.unlock();

		if (message != null) {
			connection.sendMessage(message);
		}
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.storage.AbstractProbTreeStorage#closeAll()
	 */
	@Override
	public void closeAll() {
	}
}
